-- |
-- Module:        Hexena.Math.Statistics
-- Stability:     Experimental
-- Portability:   Portable
--
-- Statistical functions
module Hexena.Math.Statistics
  ( Histogram
  , Probabilities
  , calculateE
  , calcBothHist
  , calcHist
  , calcStdDev
  , histToProbability
  , meanStdev
  , rleBS
  , scaleHist
  ) where

import Data.Array.Base (unsafeWrite, unsafeRead)
import Data.Array.IO
import Data.Array.Unboxed
import Data.IFF
import Data.List (group)
import Data.Word
import Data.Packed.Vector
import Data.Packed.Matrix

import Hexena.Math.Statistics.Internal (mean, stddev)

-- | How many times a value was found.
type Histogram = UArray Word8 Int

-- | Probability of finding a Word8 value.
type Probabilities = UArray Word8 Double

-- | Calculate the Histogram as an unboxed array. The given list of
-- Word8 is used as input.
calcHist ::  [Word8] -> IO Histogram
calcHist c = newArray (0, 255) 0 >>= m
  where m :: IOUArray Word8 Int -> IO (UArray Word8 Int)
        m a = mapM_ (f a) c >> unsafeFreeze a -- freeze a
        f a x = let x' = fromIntegral x
                in unsafeRead a x' >>= unsafeWrite a x' . (1 +)
                -- readArray a x >>= writeArray a x . (1 +)

-- | Calculate both the standard histogram and the rle-reduced
-- | histogram simultaneously.
calcBothHist :: [Word8] -> IO (Histogram, Histogram)
calcBothHist content = do
  hist <- calcHist content
  histC <- calcHist $ rleBS content
  return (hist, histC)

-- | Calculate the entropy of an histogram
calculateE :: Histogram -> Double
calculateE a = negate $ sum probs
  where total = fromIntegral $ sum $ elems a
        prob x = fromIntegral x / total
        f x = prob x * logBase 2 (prob x)
        probs = map f $ filter (0 /=) $ elems a

-- | Remove consecutive runs of duplicate characters. This removes the
-- zero byte fillings and gives a more useful estimate of the "true"
-- entropy. The difference between the two values gives us also some
-- insight into the characteristics the file.
rleBS :: [Word8] -> [Word8]
rleBS = map head . group

-- | Scale the histogram. TODO: maximum value is 256, this should be given as a variable...
scaleHist :: Histogram -> (Int -> Int) -> Histogram
scaleHist a f = amap (f . mf) a
    where m = maximum $ elems a
          mf x = x * 256 `div` m
          -- mf x = 16 * floor (logBase 2 (fromIntegral $ x * 256 `div` m))

histToProbability :: Histogram -> Probabilities
histToProbability h = listArray (0, 255) $ map f $ elems h
  where total :: Double
        total = fromIntegral $ sum $ elems h
        f = (/ total) . fromIntegral 

meanStdev :: [Double] -> (Double, Double)
meanStdev xs = (mean xs, stddev xs)

calcStdDev :: Matrix Double -> [(Double, Double)]
calcStdDev m = let mColumns = toColumns m
                   f :: Vector Double -> (Double, Double)
                   f v = (mean v', stddev v')
                     where v' = toList v
               in map f mColumns

